#include <stdint.h> 
#include <virt_mmu.h>
#include <error.h>
#include <esr.h>
#include <asm/mmpage.h>
#include <virt_io.h>

#define S2_ROOT_PG_SIZE (PAGE_SIZE * 2)
char s2_pg_dir[S2_ROOT_PG_SIZE] __attribute__((aligned(S2_ROOT_PG_SIZE)));

#define S2_IPA_BITS                    40UL

#define S2_PUD_SHIFT			30UL
#define S2_PUD_SIZE			(1UL << S2_PUD_SHIFT)
#define S2_PUD_MASK			(~(S2_PUD_SIZE - 1))

/* 注意：对于S2 concatenated页表，这里S2_PTRS_PER_PUD 和之前 略有变化*/
#define S2_PTRS_PER_PUD                 (1 << (S2_IPA_BITS - S2_PUD_SHIFT))

#define S2_PMD_SHIFT			21UL
#define S2_PMD_SIZE			(1UL << S2_PMD_SHIFT)
#define S2_PMD_MASK			(~(S2_PMD_SIZE - 1))

#define S2_PTE_SHIFT			12UL
#define S2_PTE_SIZE			(1UL << S2_PTE_SHIFT)
#define S2_PTE_MASK			(~(S2_PTE_SIZE - 1))

/* 注意：对于S2 concatenated页表，这里s2_pud_xxx 这几个宏 和之前 略有变化*/
#define s2_pud_index(addr) (((addr) >> S2_PUD_SHIFT) & (S2_PTRS_PER_PUD - 1))
#define s2_pud_offset_raw(pg_dir, addr) ((pud_t *)(((pud_t *)(pg_dir)) \
			+ s2_pud_index(addr)))

#define s2_pud_offset(pg_dir, addr) (s2_pud_offset_raw(pg_dir, (addr)))
#define __pa(x) (x)


static unsigned long get_vtcr_el2(unsigned int parange, unsigned int ipa_bits,
        unsigned int pgtable_levels, unsigned int tg){
    unsigned long value = 0;
    unsigned long sl0;

    printk("%s PARange: %d IPA bit: %d PageTable levels: %d TG: %d\n", __func__, parange, ipa_bits, pgtable_levels, tg);
    /* set T0SZ */
    value |= ((64 - ipa_bits) << 0);

    /* set tg0 */
    value |= (tg << 14);

    /* set sl0 */
	if (pgtable_levels == 4)
		sl0 = 2;
	else
		sl0 = 1;

    value |= (sl0 << 6);

	/* cacheability attribute*/
	value |= (0x1 << 8);	// Normal memory, Inner WBWA
	value |= (0x1 << 10);	// Normal memory, Outer WBWA
	value |= (0x3 << 12);	// Inner Shareable

	// PS --- pysical size 44bit
	value |= (parange << 16);
	// vmid -- 8bit
	value |= (0x0 << 19);
	printk("%s vtcr value 0x%lx\n", __func__, value);

	return value;

}

static void flush_all_vms(void)
{
	asm volatile(
		"dsb sy\n"
		"tlbi vmalls12e1is\n"
		"dsb sy\n"
		"isb\n"
		: : : "memory"
	);
}

static inline void pud_populate(pud_t *pud, unsigned long pmd_phy)
{
	set_pud(pud, __pud(pmd_phy | PMD_TYPE_TABLE));
}

static inline unsigned long pmd_alloc_one(){
    return get_free_page();
}

static int __pmd_alloc(pud_t *pud, unsigned long address){
    unsigned long pmd_phy = pmd_alloc_one();
    
    if((uintptr_t)pmd_phy == 0){
        return -ENOMEM; 
    }

	if (!pud_present(*pud))
		pud_populate(pud, pmd_phy);

    return 0;
}

static inline unsigned long pte_alloc_one(void)
{
	return get_free_page();
}


pmd_t* pmd_alloc(pud_t *pud, unsigned long address){
    return ((pud_none(*pud)) && __pmd_alloc(pud, address))?
        NULL: pmd_offset_phys(pud, address);
}

pte_t* pte_alloc(pmd_t *pmd, unsigned long address){
    return ((pmd_none(*pmd)) && __pte_alloc(pmd, address))?
        NULL: pte_offset_phys(pmd, address);
}

static inline void pmd_populate(pmd_t *pmd, unsigned long pte)
{
	set_pmd(pmd, __pmd(__pa(pte) | PMD_TYPE_TABLE));
}

int __pte_alloc(pmd_t *pmd, unsigned long address)
{
	unsigned long pte_phy = pte_alloc_one();
	if (!pte_phy)
		return -ENOMEM;

	if (!pmd_present(*pmd))
		pmd_populate(pmd, pte_phy);

	return 0;
}

static int stage2_page_fault(unsigned long gpa, unsigned long hpa, unsigned long prot)
{
    pgd_t *pgd;
    pud_t *pud;
    pmd_t *pmd;
    pte_t *pte;

	//printk("%s s2_pgdir 0x%lx\n", __func__, &s2_pg_dir);

    //pud_addr = pg_dir + pudoffset
    pud = s2_pud_offset(&s2_pg_dir, gpa);

    /*BUG: addr(pud) = 0xf4000 but be judge to NULL */
    if ((uintptr_t)pud == 0)
    {
        printk("error : pud_alloc fault \n");
        return -ENOMEM;
    }

    pmd = pmd_alloc(pud, gpa);
    if((uintptr_t)pmd == 0)
    {
        printk("error : pmd_alloc fault \n");
        return -ENOMEM;
    }

	//printk("after : pmd 0x%lx\n pmd_index 0x%x pmd_val 0x%x\n", (unsigned long)pmd, pmd_index(gpa), pmd_val(*pmd));

    pte = pte_alloc(pmd, gpa);
    if((uintptr_t)pte == 0)
        return -ENOMEM;
    
	//printk("pte 0x%lx\n pte_index 0x%x pte_val 0x%x\n", (unsigned long)pte, pte_index(gpa), pte_val(*pte));

	set_pte(pte, pfn_pte(hpa >> PAGE_SHIFT, prot));
    return 0;
}

void write_stage2_pg_reg(){
    unsigned long val;
    unsigned int parange;
    unsigned int pgtable_levels = 3;
    unsigned int ipa_bits;
    unsigned int tg = TG0_4K;

/* id_aa64mmfr0_el1 low bit */
    parange = read_sysreg(id_aa64mmfr0_el1);
    parange &= 0xf; 

    printk("%s PARange 0x%lx \n", __func__, parange);

    if(parange >= ID_AA64MMFR0_PARANGE_44) {
        parange = ID_AA64MMFR0_PARANGE_44;
        ipa_bits = 44;
        pgtable_levels = 4;
    }

    ipa_bits = 40;
    pgtable_levels = 3;

    val = get_vtcr_el2(parange, ipa_bits, pgtable_levels, tg);

    write_sysreg(val, vtcr_el2);

    write_sysreg(s2_pg_dir, vttbr_el2);

    flush_all_vms();
    return;

}

int do_vm_mem_abort(unsigned long fault_ipa, unsigned int esr, struct pt_regs *regs){
    unsigned long ec = ESR_ELx_EC(esr);
    unsigned long prot = 0;
    unsigned long far = read_sysreg(far_el2);
    unsigned long hpa;

    fault_ipa = fault_ipa << 8;
    fault_ipa &= PAGE_MASK;

    hpa = fault_ipa;

    //printk("%s fault IPA 0x%lx ec 0x%x\n", __func__, fault_ipa, ec);

    if((fault_ipa >= PBASE) && fault_ipa < PBASE_END)
        prot = S2_PAGE_DEVICE;
    else if(ec == ESR_ELx_EC_IABT_LOW)
        prot = S2_PAGE_NORMAL | S2_AP_RO;
    else
        prot = S2_PAGE_NORMAL | S2_AP_RW;

    if (check_emul_mmio_range(far))
        return emul_device(regs, far, esr);


    return stage2_page_fault(hpa, hpa, prot);

}